from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import os


if __name__ == '__main__':
    lst = []
    tsne = PCA(n_components=3)
    for i in range(300):
        np_a = np.load('data/{0}.npy'.format(i))
        lst.append(np_a)
    # for i in range(100):
    #     np_a = np.load('../data/tf/conv2d/{0}.npy'.format(i))
    #     lst.append(np_a)
    # for i in range(100):
    #     np_a = np.load('../data/torch/conv2d/{0}.npy'.format(i))
    #     lst.append(np_a)
    # for i in range(100):
    #     np_a = np.load('../data/torch/conv2d/{0}.npy'.format(i+2678))
    #     lst.append(np_a)
    np_lst = np.asarray(lst)
    np_lst = np_lst.reshape(300, 28*28)
    print(np_lst.shape)
    low_dim_data = tsne.fit_transform(np_lst)
    x_min, x_max = np.min(low_dim_data, 0), np.max(low_dim_data, 0)
    low_dim_data = low_dim_data / (x_max - x_min)
    print(low_dim_data.shape)

    # plt.title('Sample Distribution Generated By Weighted Sampling On Conv2d')
    fig = plt.figure()
    # fig.suptitle('Sample Distribution Generated By Weighted Sampling On Conv2d')
    fig.suptitle('Sample Distribution Generated By Predoo-ln1 On Conv2d')
    ax1 = Axes3D(fig)

    # x = np.linspace(-400, 400, 1000)
    # y = np.linspace(-600, 200, 1000)
    # z = np.linspace(-800, 200, 1000)
    ax1.scatter3D(low_dim_data[:, 0], low_dim_data[:, 1], low_dim_data[:, 2], cmap='Blues')

    # ax1.set_xticklabels([-0.4, -0.2, 0.0, 0.2, 0.4, 0.6], fontsize=15)
    # ax1.set_yticklabels([-0.4, -0.2, 0.0, 0.2, 0.4, 0.6], fontsize=15)
    # ax1.set_zticklabels([-0.4, -0.2, 0.0, 0.2, 0.4, 0.6], fontsize=15)
    # ax1.plot3D(x, y, z, 'gray')
    plt.show()
